"""
内容：查看深度学习基本模型驱动及其可视化功能
日期：2020年6月27日
作者：徐皓玮
"""

# 导入
from keras.utils import np_utils
from keras.models import Sequential
from keras.layers import Dense
import numpy as np
import os
import matplotlib.pyplot as plt
from keras.utils import plot_model
from keras.models import load_model

# 预设
os.environ["PATH"] += os.pathsep + 'F:/Anaconda3/Grahpviz/bin'  # 添加环境变量
LOGS_SAVE_PATH = './logs/Demo1_Acc-Loss Curve.pdf'  # 训练讯息保存路径
DATASET_PATH = '../dataset/mnist/mnist.npz'  # 数据集路径
MODEL_PATH = './infer_model/'  # 模型保存路径
MODEL_VISUAL = './logs/Model.pdf'  # 模型可视化
# 超参
N_SAMPLES_TRAIN = 60000     # 训练集样本数量
N_SAMPLES_TEST = 10000  # 测试集样本数量
IMG_SIZE = 28  # 图片尺寸
N_CLS = 10  # 类别数
BATCH_SIZE = 128  # 批次数量
EPOCH = 25  # 周期
VAL_SIZE_RATE = 0.3  # 验证集比例


def history_plot(history):
    """
    # 绘制图像
    :param history: 保存训练历史讯息的history类
    :return:
    """
    style = ['g--', 'y--', 'b-', 'r-']
    for i, key in enumerate(history.history.keys()):
        plt.plot(history.history[key], style[i])
    plt.title('Acc-Loss Curve')
    plt.ylabel('Acc/Loss')
    plt.xlabel('Epoch')
    plt.legend(['Val Loss', 'Val Accuracy', 'Loss', 'Accuracy'])
    plt.savefig(LOGS_SAVE_PATH)
    plt.show()


def load_dataset():
    """
    # 生成数据集
    :return: 划分好的训练集和测试集
    """
    with np.load(DATASET_PATH) as data:
        (X_train, Y_train), (X_test, Y_test) = (
            data['x_train'], data['y_train']), (data['x_test'], data['y_test'])
    X_train = X_train.reshape(N_SAMPLES_TRAIN,
                              IMG_SIZE * IMG_SIZE).astype('float32') / 255.0
    X_test = X_test.reshape(N_SAMPLES_TEST,
                            IMG_SIZE * IMG_SIZE).astype('float32') / 255.0
    Y_train = np_utils.to_categorical(Y_train)
    Y_test = np_utils.to_categorical(Y_test)

    print('train samples: {}\n'
          'test samples: {}'.format(X_train.shape[0], X_test.shape[0]))

    return X_train, Y_train, X_test, Y_test


def main():
    """
    # 构建并训练模型
    :return:
    """
    # 生成数据集
    X_train, Y_train, X_test, Y_test = load_dataset()
    # 构建模型
    model = Sequential()
    model.add(
        Dense(
            units=64,
            input_dim=IMG_SIZE *
            IMG_SIZE,
            activation='relu'))
    model.add(Dense(units=10, activation='softmax'))
    # 设置模型训练过程
    model.compile(
        loss='categorical_crossentropy',
        optimizer='sgd',
        metrics=['accuracy'])
    # 训练模型
    history = model.fit(
        X_train,
        Y_train,
        epochs=EPOCH,
        validation_split=VAL_SIZE_RATE,
        batch_size=BATCH_SIZE,
        verbose=1)
    # 可视化模型训练过程
    history_plot(history)
    # 评价模型
    loss_and_metrics = model.evaluate(X_test, Y_test, verbose=0)
    print('\nLoss and Metrics: {}'.format(loss_and_metrics))
    # 存储模型
    model.save(os.path.join(MODEL_PATH, 'mnist_mlp_model.h5'))
    print('Model saved successfully.')
    # 使用模型
    model = load_model(os.path.join(MODEL_PATH, 'mnist_mlp_model.h5'))
    print('Model loaded successfully.')
    plot_model(model, to_file=MODEL_VISUAL, show_shapes=True) # 深度学习模型可视化
    X_hat = X_test[0:1]
    Y_hat = model.predict(X_hat)
    print('Prediction: {}'.format(Y_hat))


if __name__ == '__main__':
    main()
